package com.ai.bayesi.model;

import com.ai.bayesi.core.BayesianNetwork;
import com.ai.bayesi.core.Node;

import java.util.*;

public class MoodPredictor {
    private BayesianNetwork network;

    // 定义所有可能的状态值
    public static final List<String> YES_NO = Arrays.asList("是", "否");
    public static final List<String> WEATHER = Arrays.asList("晴天", "雨天", "阴天");
    public static final List<String> STRESS = Arrays.asList("高", "中", "低");
    public static final List<String> SLEEP = Arrays.asList("好", "一般", "差");
    public static final List<String> SOCIAL = Arrays.asList("活跃", "一般", "不活跃");

    public MoodPredictor() {
        network = new BayesianNetwork();
        createNetwork();
    }

    private void createNetwork() {
        // 创建节点
        Node mood2 = new Node("心情", YES_NO);
        Node weather2 = new Node("天气", WEATHER);
        Node stress2 = new Node("工作压力", STRESS);
        Node sleep2 = new Node("睡眠质量", SLEEP);
        Node social2 = new Node("社交活动", SOCIAL);

        // 设置网络结构
        mood2.addParent(weather2);
        mood2.addParent(stress2);
        mood2.addParent(sleep2);
        mood2.addParent(social2);

        // 设置条件概率表(CPT)

        // 天气的先验概率
        Map<List<String>, double[]> weatherCpt = new HashMap<>();
        weatherCpt.put(Collections.emptyList(), new double[]{0.5, 0.3, 0.2}); // 晴天:0.5, 雨天:0.3, 阴天:0.2
        weather2.setCpt(weatherCpt);

        // 工作压力的先验概率
        Map<List<String>, double[]> stressCpt = new HashMap<>();
        stressCpt.put(Collections.emptyList(), new double[]{0.2, 0.5, 0.3}); // 高:0.2, 中:0.5, 低:0.3
        stress2.setCpt(stressCpt);

        // 睡眠质量的先验概率
        Map<List<String>, double[]> sleepCpt = new HashMap<>();
        sleepCpt.put(Collections.emptyList(), new double[]{0.4, 0.4, 0.2}); // 好:0.4, 一般:0.4, 差:0.2
        sleep2.setCpt(sleepCpt);

        // 社交活动的先验概率
        Map<List<String>, double[]> socialCpt = new HashMap<>();
        socialCpt.put(Collections.emptyList(), new double[]{0.3, 0.5, 0.2}); // 活跃:0.3, 一般:0.5, 不活跃:0.2
        social2.setCpt(socialCpt);

        // 心情的条件概率表
        Map<List<String>, double[]> moodCpt = new HashMap<>();

// 定义各节点的可能取值
        List<String> weatherValues = Arrays.asList("晴天", "多云", "雨天");
        List<String> stressValues = Arrays.asList("低", "中", "高");
        List<String> sleepValues = Arrays.asList("好", "一般", "差");
        List<String> socialValues = Arrays.asList("活跃", "一般", "不活跃");

// 生成所有可能的组合
        for (String weather : weatherValues) {
            for (String stress : stressValues) {
                for (String sleep : sleepValues) {
                    for (String social : socialValues) {
                        double baseProbability = 0.5;

                        // 使用乘数调整概率，避免累加导致溢出
                        if (weather.equals("晴天")) baseProbability *= 1.5;
                        if (weather.equals("雨天")) baseProbability *= 0.7;

                        if (stress.equals("低")) baseProbability *= 1.3;
                        if (stress.equals("高")) baseProbability *= 0.7;

                        if (sleep.equals("好")) baseProbability *= 1.3;
                        if (sleep.equals("差")) baseProbability *= 0.7;

                        if (social.equals("活跃")) baseProbability *= 1.3;
                        if (social.equals("不活跃")) baseProbability *= 0.7;

                        // 确保概率在有效范围内
                        baseProbability = Math.max(0.01, Math.min(0.99, baseProbability));

                        // 创建组合列表
                        List<String> combination = Arrays.asList(weather, stress, sleep, social);

                        // 设置条件概率
                        moodCpt.put(combination, new double[]{baseProbability, 1 - baseProbability});
                    }
                }
            }
        }

// 设置CPT并验证
        mood2.setCpt(moodCpt);


        // 将节点添加到网络
        network.addNode(weather2);
        network.addNode(stress2);
        network.addNode(sleep2);
        network.addNode(social2);
        network.addNode(mood2);
    }

    /**
     * 预测心情好的概率
     *
     * @param evidence 证据Map，例如 {"天气":"晴天", "工作压力":"低", ...}
     * @return 心情好的概率 [0,1]
     */
    public double predictGoodMood(Map<String, String> evidence) {
        Map<String, String> query = Collections.singletonMap("心情", "是");
        return network.infer(query, evidence);
    }

    /**
     * 获取所有可能的证据组合
     */
    public List<Map<String, String>> getAllPossibleEvidenceCombinations() {
        List<Map<String, String>> combinations = new ArrayList<>();

        for (String weather : WEATHER) {
            for (String stress : STRESS) {
                for (String sleep : SLEEP) {
                    for (String social : SOCIAL) {
                        Map<String, String> evidence = new HashMap<>();
                        evidence.put("天气", weather);
                        evidence.put("工作压力", stress);
                        evidence.put("睡眠质量", sleep);
                        evidence.put("社交活动", social);
                        combinations.add(evidence);
                    }
                }
            }
        }

        return combinations;
    }
}